rm(list = ls())
here::i_am(file.path("code", "12_fs_functional_form.R"))
library(here)
source(here("code", "config.R"))

est_df <-  read_parquet(here("data", "analysis", "analysis.parquet")) |>
  mutate(
    exemption_married = paste0(exemption, married_cards),
    order_num_serial_norm_10 = cut(order_num_serial_norm, 10),
    order_num_recenter_norm = (local_order_num_from_serial - accepted_at_camp) / registrants_second_report,
    order_num_recenter_norm_below_0 = order_num_recenter_norm < 0
  )

################################################################################
# Binscatter by married x exemption

slopes_by_type <- feols(
  vet_combined ~ order_num_serial_norm,
  split = ~ exemption_married,
  data = est_df,
  vcov = "hetero"
) |>
  sapply(\(x) x$coefficients[[2]])

p <- est_df |>
  mutate(
    exemption_married = factor(
      exemption_married,
      levels = c("FALSEFALSE", "FALSETRUE", "TRUEFALSE", "TRUETRUE"),
      labels = paste0(
        c("**No exemption claim<br>Not married**", "**No exemption claim<br>Married**",
          "**Exemption claim<br>Not married**", "**Exemption claim<br>Married**"),
        "<br><br>N = ",
        est_df |> count(exemption_married) |> pull(n) |> formatC(big.mark=","),
        "<br>Slope = ",
        sprintf("%.3f", slopes_by_type)
      )
    )
  ) |>
  ggplot(aes(x = order_num_serial_norm, y = as.integer(vet_combined))) +
  stat_binscatter() +
  my_theme() +
  labs(
    x = "Order number (scaled)",
    y = "Veteran"
  ) +
  theme(
    strip.text.x = element_markdown(face = "plain"),
  ) +
  facet_wrap(~ exemption_married, nrow = 1, scales = "free") +
  scale_x_continuous(labels = c(0, .25, .5, .75, 1))

p_color <- p +
  geom_smooth(linewidth = .5, color = "#F8766D", method = "lm", se = FALSE, formula = "y ~ x") 
ggsave(file.path(fig_dir, "color", "fs_by_type.pdf"), plot = p_color, width = 8, height = 2.5)

p_bw <- p +
  geom_smooth(linewidth = .5, color = "black", method = "lm", se = FALSE, formula = "y ~ x") 
ggsave(file.path(fig_dir, "bw", "fs_by_type.pdf"), plot = p_bw, width = 8, height = 2.5)

################################################################################
# Binscatter with order number minus # men accepted at camp (piecewise linear)

p <- ggplot(est_df, aes(x = order_num_recenter_norm, y = as.integer(vet_combined))) +
  stat_binscatter(data = est_df |> filter(order_num_recenter_norm < 0), bins = 10) +
  stat_binscatter(data = est_df |> filter(order_num_recenter_norm >= 0), bins = 10) +
  my_theme() +
  labs(
    x = "Order number minus number accepted at camp (scaled)",
    y = "Veteran"
  ) +
  xlim(-.5, .8)

p_color <- p +
  geom_smooth(
    data = est_df |> filter(order_num_recenter_norm < 0),
    linewidth = .5,
    color = "#F8766D",
    method = "lm",
    se = FALSE,
    formula = "y ~ x",
    fullrange = FALSE
  ) +
  geom_smooth(
    data = est_df |> filter(order_num_recenter_norm >= 0),
    linewidth = .5,
    color = "#F8766D",
    method = "lm",
    se = FALSE,
    formula = "y ~ x",
    fullrange = FALSE
  )
ggsave(
  file.path(fig_dir, "color", "fs_recentered.pdf"),
  plot = p_color,
  width = 5,
  height = 3
)

p_bw <- p +
  geom_smooth(
    data = est_df |> filter(order_num_recenter_norm < 0),
    linewidth = .5,
    color = "black",
    method = "lm",
    se = FALSE,
    formula = "y ~ x",
    fullrange = FALSE
  ) +
  geom_smooth(
    data = est_df |> filter(order_num_recenter_norm >= 0),
    linewidth = .5,
    color = "black",
    method = "lm",
    se = FALSE,
    formula = "y ~ x",
    fullrange = FALSE
  )
ggsave(
  file.path(fig_dir, "bw", "fs_recentered.pdf"),
  plot = p_bw,
  width = 5,
  height = 3
)

################################################################################
# Table with alternative first stage functional form

ols <- feols(
  is_naacp ~ vet_combined |
    birthyr_cards + bpl_cards + board_identifier +
    exemption^married_cards + farmer + laborer + farm_laborer,
  data = est_df,
  vcov = "hetero"
)

rhs <- model.matrix(ols)
lhs <- model.matrix(ols, type = "lhs")
my_nonvet <- mean(lhs[rhs == 0])

ols_coef <- sprintf(
  "%.04f",
  coefficients(ols)
)

ols_t <- sprintf(
  "%.02f",
  tstat(ols)
)

fs1 <- feols(
  vet_combined ~ order_num_serial_norm |
    birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer,
  data = est_df,
  combine.quick = FALSE
)

est_df$instrument1 <- predict(fs1, est_df)

# Piecewise linear
fs2 <- feols(
  vet_combined ~ order_num_recenter_norm*order_num_recenter_norm_below_0 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer,
  data = est_df,
  combine.quick = FALSE
)

est_df$instrument2 <- predict(fs2, est_df)

# Cubic polynomial
fs3 <- feols(
  vet_combined ~ order_num_serial_norm + order_num_serial_norm^2 + order_num_serial_norm^3 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer,
  data = est_df,
  combine.quick = FALSE
)

est_df$instrument3 <- predict(fs3, est_df)

# Deciles
fs4 <- feols(
  vet_combined ~ order_num_serial_norm_10 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer,
  data = est_df,
  combine.quick = FALSE
)

est_df$instrument4 <- predict(fs4, est_df)

# Deciles x exemption x married
fs5 <- feols(
  vet_combined ~ order_num_serial_norm_10*exemption_married |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer,
  data = est_df,
  combine.quick = FALSE
)

est_df$instrument5 <- predict(fs5, est_df)

res1 <- feols(
  is_naacp ~ 1 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer |
    vet_combined ~ order_num_serial_norm,
  data = est_df,
  cluster = ~ serial_num
)

res2 <- feols(
  is_naacp ~ 1 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer |
    vet_combined ~ instrument2,
  data = est_df,
  cluster = ~ serial_num
)

res3 <- feols(
  is_naacp ~ 1 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer |
    vet_combined ~ instrument3,
  data = est_df,
  cluster = ~ serial_num
)

res4 <- feols(
  is_naacp ~ 1 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer |
    vet_combined ~ instrument4,
  data = est_df,
  cluster = ~ serial_num
)

res5 <- feols(
  is_naacp ~ 1 |
    board_identifier + birthyr_cards + bpl_cards + exemption^married_cards + farmer + laborer + farm_laborer |
    vet_combined ~ instrument5,
  data = est_df,
  cluster = ~ serial_num
)

etable(
  res1,
  res2,
  res3,
  res4,
  res5,
  tex = TRUE,
  signif.code = c("***" = 0.01, "**" = 0.05, "*" = 0.10),
  fitstat = ~ n + r2 + ivwald,
  dict = c(
    "vet_combinedTRUE" = "Veteran",
    "serial_num" = "Serial num.",
    "birthyr" = "Birth year",
    "bplstr" = "Birth state",
    "board_identifier" = "Draft board",
    "farmer" = "Farmer",
    "laborer" = "Laborer",
    "farm_laborer" = "Farm laborer",
    "is_naacp" = "NAACP member",
    "statefip" = "State in 1930"
  ),
  replace = TRUE,
  file = file.path(tab_dir, "fs_robustness.tex"),
  fixef.group =  list(
    "Birth year and state" = c("Birth year", "Birth state"),
    "-_Prewar occupation" = c("Farmer", "Laborer", "Farm laborer"),
    "-_State in 1930" = c("State in 1930"),
    "-_County in 1930" = c("county1930"),
    "-_Exemption claim $\\times$ married" = c("Exemption-Married")
  ),
  postprocess.tex = function(x) { gsub("Wald (1st stage), Veteran", "First stage $F$-statistic", x, fixed = TRUE) },
  extralines = list(
    "_^Nonparametric $\\times$ exemption claim $\\times$ married" = c("", "", "", "", "\\checkmark"),
    "_^Nonparametric fit" = c("", "", "", "\\checkmark", ""),
    "_^Cubic polynomial" = c("", "", "\\checkmark", "", ""),
    "_^Piecewise linear" = c("", "\\checkmark", "", "", ""),
    "__Dep. var. mean (nonveterans)" = rep(my_nonvet, 5),
    "__OLS coefficient" = rep(ols_coef, 5),
    "__OLS $t$-statistic" = rep(ols_t, 5)
  )
)
